package com.witmore.neutron.ai.web.service;

import com.jk.common.response.ResponseResult;
import com.witmore.neutron.ai.api.dto.MessageContentDTO;
import com.witmore.neutron.ai.api.dto.MessageDTO;
import com.witmore.neutron.ai.api.enums.ChatClientType;
import com.witmore.neutron.ai.api.service.ChatClientAdapter;
import com.witmore.neutron.ai.db.entity.AgentDO;
import com.witmore.neutron.ai.db.entity.AgentSessionRecordDO;
import com.witmore.neutron.ai.db.entity.ModelConfigDO;
import com.witmore.neutron.ai.db.mapper.AgentMapper;
import com.witmore.neutron.ai.db.mapper.ModelConfigMapper;
import com.witmore.neutron.ai.web.controller.bo.MessageBO;
import com.witmore.neutron.ai.web.enums.ShowFlagEnum;
import com.witmore.neutron.ai.web.filter.GlobalAccountSession;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.springframework.beans.BeanUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.stereotype.Service;
import org.springframework.util.ObjectUtils;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.stream.Collectors;

@Slf4j
@Service
public class ChatClientRouterService {

    @Autowired
    private AgentMapper agentMapper;

    @Autowired
    private ModelConfigMapper modelConfigMapper;

    @Autowired
    private SessionRecordService sessionRecordService;


    private final Map<ChatClientType, ChatClientAdapter> chatClientAdapters;

    public ChatClientRouterService(List<ChatClientAdapter> adapters) {
        this.chatClientAdapters = adapters.stream()
                .collect(Collectors.toMap(ChatClientAdapter::getClientType, Function.identity()));
    }

    public ResponseResult<?> routeRequest(MessageBO messageBO) {
        // 1.参数校验和日志记录
        paramCheckAndLog(messageBO);
        AgentDO agentDO = agentMapper.selectById(messageBO.getAgentId());
        if (agentDO == null) {
            throw new IllegalArgumentException("Invalid request: agentId is not valid");
        }
        // 2.获取模型类型
        ModelConfigDO modelConfigDO = this.getModelConfigDO(agentDO);
        // 3.获取对应的 ChatClientAdapter
        ChatClientAdapter adapter = this.getChatClientAdapter(modelConfigDO);
        // 4.打包消息提示词
        packMessagePrompt(messageBO, agentDO);
        // 5.保存会话消息
        AgentSessionRecordDO recordDetailDO = sessionRecordService.saveSessionRecord(messageBO, agentDO);
        // 6.调用 ChatClientAdapter 生成响应
        MessageDTO messageDTO = new MessageDTO();
        BeanUtils.copyProperties(messageBO, messageDTO);
        messageDTO.setSessionId(recordDetailDO.getSessionId());
        // 只有需要登录的智能体才需要传递用户信息
        if (agentDO.getLoginFlag() != null && agentDO.getLoginFlag() == 1) {
            messageDTO.setUserPhone(GlobalAccountSession.getUserInfo() == null ? "" : GlobalAccountSession.getUserInfo().getUserPhone());
        }
        // TODO API 根据key 获取用户信息
        messageDTO.setUserId(messageBO.getUserId());
        if (messageBO.getStream()) {
            return ResponseResult.fail(500, "暂不支持sse流式接口");
        }
        return ResponseResult.success(adapter.generateResponse(modelConfigDO, messageDTO));
    }

    public Flux<ServerSentEvent<String>> routeRequestSSEStream(MessageBO messageBO) {
        // 1.参数校验和日志记录
        paramCheckAndLog(messageBO);
        AgentDO agentDO = agentMapper.selectById(messageBO.getAgentId());
        if (agentDO == null) {
            throw new IllegalArgumentException("Invalid request: agentId is not valid");
        }
        // 2.获取模型类型
        ModelConfigDO modelConfigDO = this.getModelConfigDO(agentDO);
        // 3.获取对应的 ChatClientAdapter
        ChatClientAdapter adapter = this.getChatClientAdapter(modelConfigDO);
        // 4.打包消息提示词
        packMessagePrompt(messageBO, agentDO);
        // 5.保存会话消息
        AgentSessionRecordDO recordDetailDO = sessionRecordService.saveSessionRecord(messageBO, agentDO);
        // 6.调用 ChatClientAdapter 生成响应
        MessageDTO messageDTO = new MessageDTO();
        BeanUtils.copyProperties(messageBO, messageDTO);
        messageDTO.setSessionId(recordDetailDO.getSessionId());
        // 只有需要登录的智能体才需要传递用户信息
        if (agentDO.getLoginFlag() != null && agentDO.getLoginFlag() == 1) {
            messageDTO.setUserPhone(GlobalAccountSession.getUserInfo() == null ? "" : GlobalAccountSession.getUserInfo().getUserPhone());
        }
        messageDTO.setUserId(messageBO.getUserId());
        return adapter.generateResponseStreamSSEFlux(modelConfigDO, messageDTO);
    }


    private static void packMessagePrompt(MessageBO messageBO, AgentDO agentDO) {
        // 提示词
        if ("web".equalsIgnoreCase(messageBO.getSource()) || "platform".equalsIgnoreCase(messageBO.getSource())) {
            // web端的提示词需要 库内查询组装
            messageBO.getMessages().add(new MessageContentDTO("system", agentDO.getAiSysPromptText(), ShowFlagEnum.NOT_SHOW.getCode()));
        }
    }

    @NotNull
    private ModelConfigDO getModelConfigDO(AgentDO agentDO) {
        // 获取模型类型
        ModelConfigDO modelConfigDO = modelConfigMapper.selectById(agentDO.getAiModelId());
        if (modelConfigDO == null) {
            throw new IllegalArgumentException("Invalid request: modelConfig is not valid");
        }
        // 模型未启用
        if (!modelConfigDO.getEnabled()) {
            throw new IllegalArgumentException("Invalid request: modelConfig is not valid");
        }
        return modelConfigDO;
    }

    @NotNull
    private ChatClientAdapter getChatClientAdapter(ModelConfigDO modelConfigDO) {
        // chatClient 类型
        String chatClientType = modelConfigDO.getChatClientType();
        // 转换成枚举
        ChatClientType clientType = ChatClientType.valueOf(chatClientType);
        ChatClientAdapter adapter = chatClientAdapters.get(clientType);
        if (adapter == null) {
            throw new IllegalArgumentException("Unsupported clientType: " + clientType);
        }
        return adapter;
    }

    private static void paramCheckAndLog(MessageBO messageBO) {
        // 智能助手类型必填
        if (messageBO == null || messageBO.getAgentId() == null || ObjectUtils.isEmpty(messageBO.getAgentCode())) {
            throw new IllegalArgumentException("Invalid request: agentId must be provided");
        }
        // 消息内容不能为空
        if (ObjectUtils.isEmpty(messageBO.getMessages())) {
            throw new IllegalArgumentException("Invalid request: messages must be provided");
        }
    }
}
